import random
import numpy as np
from view.loadSaveDB import loadSaveDB


class VariatePaiPartition:

    def __init__(self, data, label):
        self.data = data
        self.rows = data.shape[0]
        self.label = label
        self.cluster = len(np.unique(label))
        self.randomSampleBlocks = None
        self.randomSampleBlocks_labels = None

    def partition(self, num):
        labelIdxs = []
        for i in range(self.cluster):
            labelIdxs.append(np.where(self.label == i)[0].tolist())

        paiLst = np.linspace(0, 1, num)
        idxs = []
        for pai in paiLst:
            c1_num = int(pai * (self.rows / 2))
            c2_num = int((self.rows / 2) - c1_num)
            c1_idxs = labelIdxs[0][:c1_num]
            c2_idxs = labelIdxs[1][:c2_num]
            idxs.append(c1_idxs+c2_idxs)

        self.randomSampleBlocks = []
        self.randomSampleBlocks_labels = []
        for blockidxs in idxs:
            self.randomSampleBlocks.append(self.data[blockidxs, :])
            self.randomSampleBlocks_labels.append(self.label[blockidxs])
        self._fixLabels()

    def _fixLabels(self):
        fixedLabels = []
        for label in self.randomSampleBlocks_labels:
            fixedLabels.append(label)
            label_block = np.unique(label)
            nclusters_block = len(label_block)
            if list(label_block) == list(range(nclusters_block)):
                continue
            for i, t in enumerate(label_block):
                idxs = np.where(label == t)[0]
                label[idxs] = i
            fixedLabels[-1] = label
        self.randomSampleBlocks_labels = fixedLabels

    def saveToDB(self, des, fileName, name):
        lsd = loadSaveDB()
        ids = []
        for data, label in zip(self.randomSampleBlocks, self.randomSampleBlocks_labels):
            data = np.concatenate((label[:, np.newaxis], data), axis=1)
            id = lsd.saveDBForDataset(data=data, des=des, fileName=fileName, nClusters=max(label)+1, name=name,
                                      haveLabels=True, confirm=True)
            ids.append(id)
        return ids
